import base64
import json
import cv2
import numpy as np
from openai import OpenAI
from PIL import Image, ImageDraw, ImageFont
import os

# ==================== 1. Base64 编码函数 ====================
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

# ==================== 2. 调用视觉模型获取边界框 ====================
example_image_path = "thirdoutput/blacked_output.jpg"
original_image_path = "images/11.jpg"

base64_example = encode_image(example_image_path)
base64_original = encode_image(original_image_path)

client = OpenAI(
    api_key="sk-728ab8263bf64d5d925bc3179148549f",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)

completion = client.chat.completions.create(
    model="qwen-vl-max-latest",
    messages=[
        {
            "role": "system",
            "content": "You are a helpful assistant that performs visual reasoning across images."
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "请看这张示例图像："},
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_example}"}
                },
                {
                    "type": "text",
                    "text": "上图中被框出的多个区域是你需要关注的目标。请理解这些不同目标的语义（如内容信息，类别等）。"
                },
                {"type": "text", "text": "现在请分析下面这张图像："},
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_original}"}
                },
                {
                    "type": "text",
                    "text": (
                        "请在上图中找出所有与示例图像中被框出目标语义相同或相似度在90%的区域，可能会有多个"
                        "包括所有微小或不明显的实例。"
                        "用边界框（bbox）标注它们的位置，并描述每个目标的特征，要确保多个边界框不重合。\n"
                        "请以JSON格式输出结果，格式如下：\n"
                        "{\n"
                        "  \"objects\": [\n"
                        "    {\n"
                        "      \"bbox\": [x1, y1, x2, y2],\n"
                        "      \"description\": \"对该目标的详细描述\"\n"
                        "    }\n"
                        "  ]\n"
                        "}\n"
                        "注意：\n"
                        "- bbox 坐标为 [左上x, 左上y, 右下x, 右下y]，单位为像素\n"
                        "- 必须找出所有匹配目标，不能遗漏\n"
                        "- 不要输出任何额外文本\n"
                        "- 如果没有匹配目标，返回 {\"objects\": []}"
                    )
                }
            ]
        }
    ],
    temperature=0,
    extra_body={
        "vl_high_resolution": True,
        "top_k": 4,
        "seed": 3407
    }
)

# 获取响应
response_text = completion.choices[0].message.content.strip()

# ==================== 清理响应内容 ====================
# 去除可能存在的 Markdown 代码块标记
if response_text.startswith("```json"):
    response_text = response_text[7:]  # 去掉开头的 ```json
if response_text.endswith("```"):
    response_text = response_text[:-3]  # 去掉结尾的 ```

print("清理后的模型响应：")
print(response_text)

# 解析 JSON 结果
try:
    result = json.loads(response_text)

    # 过滤掉描述中包含“空白区域”或“内容为空”的对象
    filtered_objects = [
        obj for obj in result.get("objects", [])
        if "空白区域" not in obj.get("description", "") and "内容为空" not in obj.get("description", "")
    ]
    result["objects"] = filtered_objects

    with open("bbox_data.json", "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
except json.JSONDecodeError as e:
    print("模型未返回合法JSON：", e)
    print("原始响应内容：")
    print(response_text)
    exit(1)



# ==================== 3. 加载图像并转为 PIL 图像 ====================
image_cv = cv2.imread(original_image_path)
if image_cv is None:
    raise FileNotFoundError("图像未找到，请检查路径")

image_pil = Image.fromarray(cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB))
draw = ImageDraw.Draw(image_pil)

# ==================== 4. 加载中文字体 ====================
font_path = "C:/Windows/Fonts/simhei.ttf"  # 改为你的中文字体路径
try:
    font = ImageFont.truetype(font_path, 20)
except Exception as e:
    print("无法加载中文字体，请检查路径或安装字体。", e)
    exit(1)

# ==================== 5. 绘制边界框和中文说明 ====================
for idx, obj in enumerate(result.get("objects", []), start=1):
    x1, y1, x2, y2 = obj["bbox"]
    desc = obj.get("description", f"ROI{idx}")

    # 绘制红色矩形框
    draw.rectangle([x1, y1, x2, y2], outline="red", width=2)

    # 写中文说明
    text_pos = (x1, y1 - 25 if y1 - 25 > 5 else y1 + 5)
    draw.text(text_pos, desc, font=font, fill="red")

# ==================== 6. 转回 OpenCV 图像格式并保存到 thirdoutput 目录 ====================
os.makedirs("thirdoutput", exist_ok=True)  # 确保目录存在
result_img = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
output_image_path = "thirdoutput/labeled_with_chinese2.jpg"
cv2.imwrite(output_image_path, result_img)

print(f"标注完成，结果图片已保存至 {output_image_path}")